import concurrent.futures
import json
import re

import loguru
import pandas as pd
import requests
from tqdm import tqdm

from tools import get_mysql_con


class SparkAiApiHttp:
    def __init__(self,
                 auth_text: str = None,
                 spark_api_secret: str = None,
                 spark_api_key: str = None
                 ):
        """
        初始化
        :param auth_text:身份
        """
        self.spark_api_secret = spark_api_secret or 'N2I2YTQ1ZTJhMTFkZmQ1MTQwMWY5MzVj',
        self.spark_api_key = spark_api_key or '98a19bd641316261ca76b50f684f5fbc',
        self.headers = {
            'Content-Type': 'application/json',
            'Authorization': f'Bearer {self.spark_api_key}:{self.spark_api_secret}'
        }
        self.auth_text = auth_text or '你是高分子高分子材料领域塑化行业物性信息高级审核员'
        
        self.url = 'https://spark-api-open.xf-yun.com/v1/chat/completions'
        self.pool = {get_mysql_con(host='192.168.101.24', database='python_etl_temp') for _ in range(2)}
    
    def get_cli(self):
        return self.pool.pop()
    
    def rep(self, data):
        self.pool.add(data)
    
    def get_data(self, text: str) -> dict:
        """

        :param text: 问题
        :return:
        """
        data = {
            "model": "general",  # 指定请求的模型
            "messages": [
                {
                    "role": "system",
                    "content": self.auth_text
                },
                {
                    "role": "user",
                    "content": text
                    
                }
            ],
        }
        return data
    
    @staticmethod
    def get_json(data):
        data = re.findall('\{.+}', data, re.S)
        if data:
            return data[0].replace('\'', '"')
    
    def save(self, save_data):
        con, cursor = self.get_cli()
        sql = 'insert into auto_project2(source_name,gpt_data) values(%s,%s)'
        cursor.execute(sql, save_data)
        con.commit()
        self.rep((con, cursor))
        loguru.logger.info(f'存储{save_data}')
        return True
    
    def request(self, source):
        """
        :param text: 问题
        :return:
        """
        text = f"""
              你是高分子高分子材料领域塑化行业的制成品实体高级审核员
              我现在给你一个文本，你告诉我这个文本是不是一个制品，因为我的文本不是一个标准名称，需要你提取出来制品的名称。
              【示例输出】：""" + """{"制品名称":"","是否是制品":""}
              【提示】：是否是制品字段只能有【是】或者【否】两中
              """ + f"需要处理的文本:{source}"
        data = self.get_data(text=text)
        try:
            response = requests.post(self.url, headers=self.headers, json=data)
            if response.status_code != 200:
                # 请求失败
                loguru.logger.info('请求code不为200')
                return
            result = response.json()['choices'][0]['message']['content']
            result = self.get_json(result)
            if not result:
                # 没有接送数据
                result = json.dumps({"制品名称": source, "是否是制品": "否"}, ensure_ascii=False)
            result = json.loads(result)
        except Exception as e:
            loguru.logger.error(f'请求错误：{e}')
            return
        self.save([source, json.dumps(result, ensure_ascii=False)])
    
    def get_source_data(self) -> list:
        """
        获取数据，
        :return:
        """
        con, cursor = get_mysql_con(host='192.168.101.24', database='python_etl_load')
        data = set(pd.read_sql('select product_label from dim_product', con)['product_label'].values.tolist())
        con.close()
        
        con, cursor = self.get_cli()
        # source1 = set(pd.read_sql('select source_name from auto_project', con)['source_name'].values.tolist())
        source2 = set(pd.read_sql('select source_name from auto_project2', con)['source_name'].values.tolist())
        self.rep((con, cursor))
        
        return list(data - source2)


def main():
    spark_obj = SparkAiApiHttp()
    data = spark_obj.get_source_data()
    with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
        futures = [executor.submit(spark_obj.request,
                                   i
                                   ) for i in data]
        for future in tqdm(concurrent.futures.as_completed(futures), total=len(futures)):
            try:
                future.result()  # 确保线程执行完成
            except Exception as e:
                print(f"Thread generated an exception: {e}")


if __name__ == '__main__':
    main()

